# first:feature_extraction
# second:classification
import os

import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from torchvision import transforms
from torchvision import datasets
import torch.nn.functional as F
import torch.optim as optim
# padding:填充附近
# bias：偏置
# stride:步长
transform=transforms.ToTensor()
train_set=datasets.MNIST(root='./data',train=True,transform=transform,download=True)
test_set=datasets.MNIST('./data',train=False,transform=transform,download=True)

train_loader=DataLoader(dataset=train_set,batch_size=50,shuffle=True)
test_loader=DataLoader(dataset=test_set,batch_size=32,shuffle=True)

class ConNet(nn.Module):
    def __init__(self):
        super(ConNet, self).__init__()
        self.conv1=nn.Conv2d(1,10,kernel_size=5)
        self.conv2=nn.Conv2d(10,20,kernel_size=5)
        self.pooling=nn.MaxPool2d(2)
        self.fc=nn.Linear(320,10)

    def forward(self,x):
        batch_size=x.size(0)
        x=F.relu(self.pooling(self.conv1(x)))
        x=F.relu(self.pooling(self.conv2(x)))
        x=x.view(batch_size,-1)
        x=self.fc(x)
        return x

model=ConNet()

criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model.to(device)
print(f'使用设备:{device}')

def train(epoch):
    running_loss=0.0
    for batch_idx,data in enumerate(train_loader,0):
        input,target=data
        input,target=input.to(device),target.to(device)
        optimizer.zero_grad()
        outputs=model(input)
        loss=criterion(outputs,target)
        loss.backward()
        optimizer.step()
        running_loss+=loss.item()
        if batch_idx % 300==299:
            print('[%d, %5d] loss: %.3f' %
                  (epoch + 1, batch_idx + 1, running_loss / 2000))
            running_loss = 0.0

def test():
    corret=0
    total=0
    with torch.no_grad():
        for data in train_loader:
            input,target=data
            input,target=input.to(device),target.to(device)
            outputs=model(input)
            _, predict=torch.max(outputs.data,dim=1)
            total+=target.size(0)
            corret+=(predict==target).sum().item()
    print('Accuracy on test set: %d %% [%d/%d]'% (100*corret/total,corret,total))

if __name__=='__main__':
    for epoch in range(5):
        train(epoch)
    test()